// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/testing/functions/prestosql/function_base_test.h>

using namespace kumo::pollux;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::functions::test;

namespace {

class CardinalityTest : public FunctionBaseTest {
 protected:
  void testArrayCardinality(
      std::function<vector_size_t(vector_size_t /* row */)> sizeAt,
      std::function<bool(vector_size_t /*row */)> is_null_at) {
    vector_size_t numRows = 100;
    auto array_vector =
        make_array_vector<int64_t>(numRows, sizeAt, value_at, is_null_at);
    auto result = evaluate<SimpleVector<int64_t>>(
        "cardinality(c0)", make_row_vector({array_vector}));
    for (vector_size_t i = 0; i < numRows; ++i) {
      EXPECT_EQ(result->is_null_at(i), array_vector->is_null_at(i)) << "at " << i;
      if (!array_vector->is_null_at(i)) {
        EXPECT_EQ(result->value_at(i), sizeAt(i)) << "at " << i;
      }
    }
  }

  void testMapCardinality(
      std::function<vector_size_t(vector_size_t /* row */)> sizeAt,
      std::function<bool(vector_size_t /*row */)> is_null_at) {
    vector_size_t numRows = 100;
    auto map_vector = make_map_vector<int64_t, int64_t>(
        numRows, sizeAt, value_at, value_at, is_null_at);
    auto result = evaluate<SimpleVector<int64_t>>(
        "cardinality(c0)", make_row_vector({map_vector}));
    for (vector_size_t i = 0; i < numRows; ++i) {
      EXPECT_EQ(result->is_null_at(i), map_vector->is_null_at(i)) << "at " << i;
      if (!map_vector->is_null_at(i)) {
        EXPECT_EQ(result->value_at(i), sizeAt(i)) << "at " << i;
      }
    }
  }

  static inline vector_size_t value_at(vector_size_t idx) {
    return idx + 1;
  }
};
} // namespace

TEST_F(CardinalityTest, array) {
  auto sizeAt = [](vector_size_t row) { return 1 + row % 7; };
  testArrayCardinality(sizeAt, nullptr);
  testArrayCardinality(sizeAt, null_every(5));
}

TEST_F(CardinalityTest, map) {
  auto sizeAt = [](vector_size_t row) { return 1 + row % 7; };
  testMapCardinality(sizeAt, nullptr);
  testMapCardinality(sizeAt, null_every(5));
}
